package com.gnote.langchain.controller;

import com.gnote.langchain.config.AiConfig;
import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.community.model.dashscope.QwenStreamingChatModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import dev.langchain4j.service.TokenStream;
import lombok.RequiredArgsConstructor;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

/**
 * @author 经纬
 * @since 2025/3/27 14:59
 */
@RestController
@RequiredArgsConstructor
public class ChatLanguageModelController {
    private final QwenChatModel chatModel;
    private final QwenStreamingChatModel streamingChatModel;
    private final AiConfig.AssistantUnique assistantUnique;


    @GetMapping("/chat")
    public String model(@RequestParam(defaultValue = "你是谁") String message) {
        return chatModel.chat(message);
    }


    @GetMapping(value = "/chat/stream", produces = "text/stream;charset=UTF-8")
    public Flux<String> stream(@RequestParam(defaultValue = "你是谁") String message) {
        return Flux.create(fluxSink -> streamingChatModel.chat(message, new StreamingChatResponseHandler() {
            @Override
            public void onPartialResponse(String s) {
                fluxSink.next(s);
            }

            @Override
            public void onCompleteResponse(ChatResponse chatResponse) {
                fluxSink.complete();
            }

            @Override
            public void onError(Throwable throwable) {
                fluxSink.error(throwable);
            }
        }));
    }


    /**
     * 记忆分割
     */
    @GetMapping(value = "/chat/stream/split", produces = "text/stream;charset=UTF-8")
    public Flux<String> split(@RequestParam String message,
                              @RequestParam Integer memoryId) {

        TokenStream stream = assistantUnique.stream(memoryId, message);
        return Flux.create(sink -> stream.onCompleteResponse(c -> sink.complete())
                .onPartialResponse(sink::next)
                .onError(sink::error)
                .start());
    }


}